package com.cetc.sdp.kmga.cs.stream.nv;

import com.rabbitmq.client.QueueingConsumer;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.streaming.api.java.JavaDStream;
import org.apache.spark.streaming.api.java.JavaStreamingContext;
import org.apache.spark.streaming.rabbitmq.RabbitMQUtils;
import org.apache.spark.streaming.rabbitmq.distributed.JavaRabbitMQDistributedKey;
import org.apache.spark.streaming.rabbitmq.models.ExchangeAndRouting;

import java.util.*;

/**
 * @description:
 * @author： DengQiang
 * @date: 2018/4/3 10:14
 */
public class RabbitMqDStream extends AbstractDStream {

    private List<String> queue;

    private Map<String, String> rabbitmqParams;

    private String exchangeName = "";

    public RabbitMqDStream(String queue, Map<String, String> rabbitmqParams) {
        this(Collections.singletonList(queue), rabbitmqParams);
    }

    public RabbitMqDStream(List<String> queue, Map<String, String> rabbitmqParams) {
        this.queue = new ArrayList<>();
        this.queue.addAll(queue);
        this.rabbitmqParams = rabbitmqParams;
    }

    public void setExchangeName(String exchangeName) {
        this.exchangeName = exchangeName;
    }

    @Override
    public JavaDStream<String> getDStreamAsPlaintext(JavaStreamingContext jssc) {
        List<JavaRabbitMQDistributedKey> distributedKeys = new LinkedList<>();
        this.queue.forEach(q -> distributedKeys.add(new JavaRabbitMQDistributedKey(q, new ExchangeAndRouting(exchangeName), rabbitmqParams)));

        Function<QueueingConsumer.Delivery, String> messageHandler = (Function<QueueingConsumer.Delivery, String>) message -> new String(message.getBody());
        return RabbitMQUtils.createJavaDistributedStream(jssc, String.class, distributedKeys, rabbitmqParams, messageHandler);
    }
}
